{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### AutoGen Introduction\n",
    "\n",
    "AutoGen AgentChat!\n",
    "\n",
    "This is a simple file to introduce me to [AutoGen(AG)](https://microsoft.github.io/autogen/stable//index.html), specifically AG AgentChat, a framework for building AI Agents and applications developed by Microsoft."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing the necessary libraries\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading the environment variables\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using Google's Gemini Models.\n",
    "Noted that Gemini models are currently in experimental state from the decs, this means that not all models are available for used in Autogen AgentsChat, especially the recently released models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gemini_client = OpenAIChatCompletionClient(\n",
    "    model=\"gemini-2.0-flash\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Applying the message to model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.messages import TextMessage\n",
    "message = TextMessage(content=\"I'd like to go to Mombasa\", source=\"user\")\n",
    "message"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Introducing the Agent using AutoGen AgentsChat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.agents import AssistantAgent\n",
    "\n",
    "travel_agent = AssistantAgent(\n",
    "    name=\"travel_assistant_agent\",\n",
    "    model_client=gemini_client,\n",
    "    system_message=\"You are a helpful assistant for an airline. You give short, clear answers.\",\n",
    "    model_client_stream=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Putting everything together with on_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_core import CancellationToken\n",
    "\n",
    "response = await travel_agent.on_messages([message], cancellation_token=CancellationToken())\n",
    "response.chat_message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Let's make a local database of ticket prices. This will act as a long term memory where the agent can refer to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sqlite3\n",
    "\n",
    "# Delete DB if it exists\n",
    "if os.path.exists(\"tickets.db\"):\n",
    "    os.remove(\"tickets.db\")\n",
    "\n",
    "# Create the DB and table\n",
    "connection = sqlite3.connect(\"tickets.db\")\n",
    "conn = connection.cursor()\n",
    "conn.execute(\"CREATE TABLE cities (city_name TEXT PRIMARY KEY, round_trip_price REAL)\")\n",
    "connection.commit()\n",
    "connection.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Populating db\n",
    "def save_city_price(city_name, round_trip_price):\n",
    "    connection = sqlite3.connect(\"tickets.db\")\n",
    "    conn = connection.cursor()\n",
    "    conn.execute(\"REPLACE INTO cities (city_name, round_trip_price) VALUES (?, ?)\", (city_name.lower(), round_trip_price))\n",
    "    connection.commit()\n",
    "    print(f\" City {city_name} with round trip price {round_trip_price} added to db.\")\n",
    "    connection.close()\n",
    "\n",
    "save_city_price(\"Nairobi\", 500)\n",
    "save_city_price(\"Nakuru\", 350)\n",
    "save_city_price(\"Eldoret\", 200)\n",
    "save_city_price(\"Kisumu\", 180)\n",
    "save_city_price(\"Mombasa\", 290)\n",
    "save_city_price(\"Narok\", 420)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to get the price of a city\n",
    "\n",
    "def get_city_price(city_name: str) -> float | None:\n",
    "    \"\"\" Getting the round trip price to travel to a city in Kenya \"\"\"\n",
    "    connection = sqlite3.connect(\"tickets.db\")\n",
    "    conn = connection.cursor()\n",
    "    conn.execute(\"SELECT round_trip_price FROM cities WHERE city_name = ?\", (city_name.lower(),))\n",
    "    result = conn.fetchone()\n",
    "    connection.close()\n",
    "    return result[0] if result else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_city_price(\"Narok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Now initilizing an agent that has access to tools, specifically the get city price"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.agents import AssistantAgent\n",
    "\n",
    "travel_assistant_agent = AssistantAgent(\n",
    "    name=\"travel_assistant_agent\",\n",
    "    model_client=gemini_client,\n",
    "    system_message=\"You are a helpful assistant for an airline. You give short, clear answers, including the price of a roundtrip ticket.\",\n",
    "    model_client_stream=True,\n",
    "    tools=[get_city_price],\n",
    "    reflect_on_tool_use=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Invoking the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await travel_assistant_agent.on_messages([message], cancellation_token=CancellationToken())\n",
    "\n",
    "for inner_message in response.inner_messages:\n",
    "    print(inner_message.content)\n",
    "response.chat_message.content"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
